from model.resnet import ResNet1D
import torch
import os
import wav2clip
import librosa
import numpy as np
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. 创建模型
model = ResNet1D(num_class=2).to(device)

# 1.1 如果存在模型 则加载
if os.path.exists("./model/save/resnet.pth"):
    model.load_state_dict(torch.load("./model/save/resnet.pth"))

# 1.2 加载测试数据集
audio, sample_rate = librosa.load("./test_data/A_wav.wav", sr=None)
inputs_press = np.loadtxt("./test_data/B_press.txt").reshape(1, 1, 300)


# 2. 加载测试数据
audio_model = wav2clip.get_model()
audio_embed = wav2clip.embed_audio(audio, audio_model).reshape(1, 1, -1)

# 2.1 转换为类型
inputs_press, audio_embed = torch.from_numpy(inputs_press).to(device).float(), torch.from_numpy(audio_embed).to(device).float()

output = model(inputs_press, audio_embed)
probabilities = F.softmax(output, dim=1)
print("二分类的概率：", probabilities)
print("测试结果标签： ", torch.argmax(probabilities, dim=1).item())

